import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.HashSet;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import org.jgrapht.graph.*;
import org.jgrapht.traverse.GraphIterator;
import org.jgrapht.traverse.TopologicalOrderIterator;
import org.jgrapht.alg.TransitiveClosure;

/**
 * Manages the whole GO DAG.
 * @author odedmagg
 *
 */
public class AnnotationsManager {
	private SimpleDirectedGraph<String, DefaultEdge> _GOdag = null;
	private SimpleDirectedGraph<String, DefaultEdge> _transitiveClosure = null;
	private int _sampleSpaceSize;
	private HashMap<String, HashSet<String>> _counts;
	
	// TEMP TODO IMPORTANT BUG
	// Currentlly, if some gene has several annotation with the same parent, that parent is counted several times!!
	// Need to Fix
	
	public int getSampleSpaceSize() {
		return _sampleSpaceSize;
	}
	
	public void setSampleSpaceSize(int size) {
		_sampleSpaceSize = size;
	}
	
	/**
	 * Loads the GO ontology and associations from file.
	 * @param ontologyFilename Name of a file containing a legal GO ontology.
	 * @param associationFilename Name of a file containing a legal GO association.
	 * @param geneMap The genes in the current network.
	 * @return Amount of genes with annotations in the genome.
	 */
	public int load(String ontologyFilename, String associationFilename, Map<String, Gene> geneMap) {
		loadOntology(ontologyFilename);
		return loadAssociation(associationFilename, geneMap);
	}
	
	@SuppressWarnings("unchecked")
	private void loadOntology(String ontologyFilename) {
		
		// TEMP TODO consider dividing to functions.
		_GOdag = new  SimpleDirectedGraph<String, DefaultEdge>(DefaultEdge.class);
		_counts = new HashMap<String, HashSet<String>>();
		
		// Load the ontology file and GO hierarchy.
		BufferedReader read;
		String currentAnnotation = null;
		String line = null;
		String parentAnnotation;
		int annotStartInd;
		int annotEndInd;
		boolean isTerm = false;

		try {
				read = new BufferedReader(new FileReader(ontologyFilename));
				while ((line = read.readLine()) != null) {
					if (line.matches("^\\[.*\\]")) {
						currentAnnotation = null;
						// System.out.println("New section"); // TEMP TODO remove
						isTerm = line.indexOf("Term") != -1;
					}
					else if (isTerm) {
						if (line.matches("^id:.*")) {
	
							// If id: tag - the current annotation is taken (assuming GO:ddddddd format).
							// Actual assumption: annotation starts with 'G'.
							annotStartInd = line.indexOf('G');
							currentAnnotation = line.substring(annotStartInd);
							_GOdag.addVertex(currentAnnotation);
							_counts.put(currentAnnotation, new HashSet<String>());

							// System.out.println("id: " + currentAnnotation); // TEMP TODO remove
						}
						else if (line.matches("^is_a:.*") && (currentAnnotation != null)) {
							// if is_a - add an edge between current annotation and its direct ancestor.
							// Assumes that the structure is is_a: GO:ddddddd
							annotStartInd = line.indexOf('G');
							annotEndInd  = line.indexOf(' ', annotStartInd);
							parentAnnotation = line.substring(annotStartInd, annotEndInd);

							if (!_GOdag.containsVertex(parentAnnotation))
								_GOdag.addVertex(parentAnnotation);
							
							_GOdag.addEdge(currentAnnotation, parentAnnotation);
							// System.out.println(currentAnnotation + " is a " + parentAnnotation);
						}
					}
				}			
				read.close();
		} catch (IOException ioe) {
			// TEMP TODO take care.
		}
		
		// TEMP TODO remove
		// System.out.println("No of vertices: " + _GOdag.vertexSet().size());
		// System.out.println("No of edges: "   + _GOdag.edgeSet().size());

		/*		// TEMP TODO remove
		CycleDetector<String, DefaultEdge> cycleDetector = new CycleDetector<String, DefaultEdge>(_GOdag);
		if (cycleDetector.detectCycles()) {
			System.out.println("Has cycles!");
		} else {
			System.out.println("Does not have cycles!");			
		}
		*/
		// TEMP TODO end remove
		
		_transitiveClosure = (SimpleDirectedGraph<String, DefaultEdge>)_GOdag.clone();
		TransitiveClosure.INSTANCE.closeSimpleDirectedGraph(_transitiveClosure);

		// TEMP TODO remove
		/*
		if (_GOdag.containsEdge("GO:0009712","GO:0018958")) {
			System.out.println("GO DAG Has direct edge\tGOOD");
		} else {
			System.out.println("GO DAG hasnt direct edge\tBAD");			
		}
		
		if (_GOdag.containsEdge("GO:0009712","GO:0009987")) {
			System.out.println("GO DAG Has path as edge\tBAD");
		} else {
			System.out.println("GO DAG hasnt path as edge\tGOOD");			
		}

		if (_transitiveClosure.containsEdge("GO:0009712","GO:0018958")) {
			System.out.println("Closure Graph Has direct edge\tGOOD");
		} else {
			System.out.println("Closure Graph hasnt direct edge\tBAD");			
		}
		
		if (_transitiveClosure.containsEdge("GO:0009712","GO:0009987")) {
			System.out.println("Closure Graph Has path as edge\tGOOD");
		} else {
			System.out.println("Closure Graph hasnt path as edge\tBAD");			
		}		
		*/
	}
	
	
	int loadAssociation(String associationFilename, Map<String, Gene> geneMap) {
		HashSet<String> genesInGenome = new HashSet<String>();
		String[] data;
		String annot;
		String prot;
		BufferedReader read;
		String line = null;
		
		if (_GOdag == null) 
			return -1;
			// TEMP TODO consider error message.
		
		// Load association file.
		try {
			read = new BufferedReader(new FileReader(associationFilename));
			while ((line = read.readLine()) != null) {
				 data  = line.split("\t");
				 prot   = data[2];
				 annot = data[3];
				 if (! annot.matches("^GO.*") ) {
					 annot = data[4];
				 }
				 
				 // If the annotation does not appear in the graph - skip it.
				 if (!_GOdag.containsVertex(annot))
					 continue;
				 
				 genesInGenome.add(prot);
				 
				 // Add the current annotation to its related gene.
				 Gene currGene = geneMap.get(prot); 
				 // If gene exists in network, add annotation.
				 if (currGene != null) {
					 currGene.addAnnotation(annot);
				 }
				 
				 // Add current entry to the genes count of the current annotation and its ancestors.
				 HashSet<String> annotGenes = _counts.get(annot);
				 	 
				 try { // TEMP TODO
					 
					 annotGenes.add(prot); // The only to keep.					 
					 
				 }
				 catch (NullPointerException npe) {
					 System.out.println("Annotation: " + annot);
					 System.out.println("Gene: " + prot);					 
					 System.exit(0);
				 } // End TEMP TODO

				for (String ancestor: ancestorsOf(annot)) {
					annotGenes = _counts.get(ancestor);
					annotGenes.add(prot);
				}

			}
			read.close();
		} catch (IOException ioe) {
			System.out.println("Can't open "+associationFilename);
		}		
		
		_sampleSpaceSize = genesInGenome.size();
		return _sampleSpaceSize;
	}
	
	/**
	 * Checks if a gene is annotated with some annotation.
	 * A gene is considered to be annotated with an annotation if the annotation was mapped to it
	 * in the association file, or if the annotation is an ancestor of another annotation that was mapped
	 * to the gene in the association file.
	 * @param gene 
	 * @param annotation
	 * @return true iff gene is annotated with annotation.
	 */
	public boolean isGeneAnnotated(Gene gene, String annotation) {
		if ((gene == null) || (annotation == null))
			return false;

		String currAnnot;
		
		Iterator<String> annotIter = gene.annotationsIterator();
		while(annotIter.hasNext()) {
			currAnnot = annotIter.next();
			if (currAnnot.equals(annotation)) {
				return true;
			}
			if (_transitiveClosure.containsEdge(currAnnot, annotation)) {
				return true;
			}
		}
		
		return false; 
	}

	/**
	 * Checks if the manager is familiar with an annotation.
	 * @param annotation the annotation to look for.
	 * @return true iff annotation is known to the annotations manager.
	 */
	public boolean containsAnnotation(String annotation) {
		return _GOdag.containsVertex(annotation);
	}
	
	/**
	 * @return A toplogical-order iterator over the entire DAG.
	 */
	public GraphIterator<String, DefaultEdge> toplogicalIterator() {
		return new TopologicalOrderIterator<String, DefaultEdge>(_GOdag);
	}
	
	/* Ancestors - related methods */
	public boolean isParentOf(String annotation, String parent) {
		return _GOdag.containsEdge(annotation, parent);
	}
	
	/**
	 * Check if annotation ancestor is an ancestor of annotation son. 
	 * @param annotation Annotation.
	 * @param ancestor potential ancestor of the 'son' annotation.
	 * @return true iff ancestor is an ancestor of son.
	 */
	public boolean isAncestorOf(String annotation, String ancestor) {
		return _transitiveClosure.containsEdge(annotation, ancestor);
	}	
	
	/**
	 * Returns the amount of ancestors of the given annotation.
	 * @param annotation
	 * @return
	 */
	public int countAncestors(String annotation) {
		return countOutgoingEdges(_transitiveClosure, annotation);
	}
	
	public int countParents(String annotation) {
		return countOutgoingEdges(_GOdag, annotation);
	}
	
	private int countOutgoingEdges(SimpleDirectedGraph<String, DefaultEdge> graph, String annotation) {
		if (!graph.containsVertex(annotation))
			return -1;
		
		return graph.outDegreeOf(annotation);		
	}
	
	/**
	 * Returns the set of ancestors of the given annotation.
	 * A vertex is an ancestor of annotation if there's a path from the annotation to the ancestor.
	 * @param annotation
	 * @return a set with all of the annotation's ancestors, or null if the annotation does not exist.
	 */
	public HashSet<String> ancestorsOf(String annotation) {
		return directlyReachableVertices(_transitiveClosure, annotation);
	}
	
	/**
	 * Returns the set of direct parents of the given annotation.
	 * Parents are vertices on connected to the annotation by an outgoing edge (annotation -> parent).
	 * @param annotation
	 * @return a set with all of the annotation's ancestors, or null if the annotation does not exist.
	 */
	public HashSet<String> parentsOf(String annotation) {
		return directlyReachableVertices(_GOdag, annotation);
	}
		
	private HashSet<String> directlyReachableVertices(
			SimpleDirectedGraph<String, DefaultEdge> graph, String annotation) {		
		if (!graph.containsVertex(annotation))
			return null;
		
		HashSet<String> result = new HashSet<String>(graph.outDegreeOf(annotation));
		for (DefaultEdge edge: graph.outgoingEdgesOf(annotation)) {
			result.add(graph.getEdgeTarget(edge));
		}
		
		return result;
	}
	
	/* Annotation counts - related methods. */
	
	/**
	 * Returns the amount of genes that were specifically annotated with this annotation.
	 * @return 0 if the annotation did not appear in the annotations file, or the amount of times it did appear. 
	 */
	public int count(String annotaion) {
		if (!_GOdag.containsVertex(annotaion))
			return 0;
		
		return _counts.get(annotaion).size();
	}
	
	/**
	 * Returns the probability / frequency of genes in the genome that are annotation with
	 * the given annotation.
	 * @param anotation
	 * @return a value between 0 and 1 indicated the fraction of annotated genes that have this annotation. 
	 */
	public double prob(String annotation) {
	   
	    double prob = 1.0 * count(annotation) / _sampleSpaceSize;
		//System.out.println("count:"+count(annotation)+"\tprob:"+prob);
		return prob;
	}
	
	
	/**
	 * returns the ancestor with the minimal probability
	 * among the common ancestors.
	 * @param commonAncestors the common ancestors
	 * @return the ancestor with the minimal probability
	 */
	private double getMinAnnot(HashSet<String> commonAncestors) {
		double min = 100000000.0;
		for (String ancestor: commonAncestors) {
			if (prob(ancestor) < min) {
				min = prob(ancestor);
			}
		}
		return min;
	}
	
	
	/**
	 * Finds the maximal similarity in annotations between gene1 and gene2.
	 * @param gene1 the first gene.
	 * @param gene2 the second gene.
	 * @return the maximal similarity between gene1 and gene2.
	 */
	public double getAnnotSimilarity(Gene gene1, Gene gene2) {
		String annot1, annot2;
		HashSet<String> commonAncestors;
		// probability of the minimum ancestor.
		double minAncestor;
		double similarity;
		double minProb = 2;//Initialize to a high value.
		// If one of the genes ins'e annotated return 0.
		if (!gene1.isAnnotated() || !gene2.isAnnotated()) {
			return 0;
		}
		Iterator<String> gene1AnnotIter = gene1.annotationsIterator();
		Iterator<String> gene2AnnotIter = gene2.annotationsIterator();
		while (gene1AnnotIter.hasNext()) {
			annot1 = gene1AnnotIter.next();
			while (gene2AnnotIter.hasNext()) {
				annot2 = gene2AnnotIter.next();
				if (annot1.equals(annot2)) {
					minAncestor = prob(annot1);
				} else if (isAncestorOf(annot1, annot2)) {
					minAncestor = prob(annot2);
				} else if (isAncestorOf(annot2, annot1)) {
					minAncestor = prob(annot1);
				} else {
					commonAncestors = getCommonAncestors(annot1, annot2);
					if (commonAncestors.size()==0) {
						minAncestor = 1;// as a result the similarity will be 0.
					} else {
						minAncestor = getMinAnnot(commonAncestors);
					}
				}
				if (minAncestor < minProb) {
				    minProb = minAncestor;
				}
			}
		}
		similarity = 1.0-minProb;
		//System.out.println(gene1.getName()+"\t"+gene2.getName()+"\t"+minProb+"\t"+similarity);
		return similarity;
	}
	
	/**
	 * @param annot1 - annotation of the first gene
	 * @param annot2 - annotation of the second gene.
	 * @return all the common ancestors of annot1 and annot2 or null if they are not valid annotations.
	 * 
	 */
	private HashSet<String> getCommonAncestors(String annot1, String annot2) {
		HashSet<String> annot1Ance = ancestorsOf(annot1);
		HashSet<String> annot2Ance = ancestorsOf(annot2);
		
		if ((annot1Ance == null) || (annot2Ance == null)) {
			return null;
		}
		
		HashSet<String> intersection = (HashSet<String>)annot1Ance. clone();
		intersection.retainAll(annot2Ance);
		
		return intersection;
	}
}

